#!/usr/bin/python
# -*-coding:utf-8-*-

'''财务数据挖掘专用算子'''

import os
import numpy as np
import pandas as pd
from joblib import wrap_non_picklable_objects

import talib as talib

from gplearn_evolve.base.utils import d_type_exam
# from gplearn_evolve.base.utils import d_map
from gplearn_evolve.base.utils import d_map_limit

from gplearn_evolve.base.utils import keys_core_pd
from gplearn_evolve.base.utils import keys_core_pd_v2

from gplearn_evolve.base.config import cw_low_limit
from gplearn_evolve.base.config import cw_up_limit

# 1
@wrap_non_picklable_objects
def def_keyfunc_ts_cw_ma(keys, key_by_path=True):
    def gp_ts_cw_ma(X, d):
        if key_by_path:
            # load key
            load_keys = np.load(keys,
                                mmap_mode='r')
        else:
            load_keys = keys.copy()

        if len(load_keys) != len(X):
            return X

        # params = d_map(d_type_exam(d), base=base)
        params = d_map_limit(d_type_exam(d), base=2, lowlimit=cw_low_limit, uplimit=cw_up_limit)

        if not key_by_path:
            values = keys_core_pd(keys, X)
        else:
            values = keys_core_pd_v2(load_keys, X)

        values = values.groupby(as_index=False,
                                level='stock_code',
                                sort=False).apply(lambda x:x.rolling(window=params,
                                                                       min_periods=params).mean())

        # TODO - 将inf设置为nan
        values[np.isinf(values['f1'])] = np.nan

        del load_keys

        return values['f1'].values

    return gp_ts_cw_ma

# 2
@wrap_non_picklable_objects
def def_keyfunc_cw_change_rate_qoq(keys, key_by_path=True):
    def gp_cw_change_rate_qoq(X):
        if key_by_path:
            # load key
            load_keys = np.load(keys,
                                mmap_mode='r')
        else:
            load_keys = keys.copy()

        if len(load_keys) != len(X):
            return X

        # params = d_map(d_type_exam(d), base=base)
        params = 1

        if not key_by_path:
            values = keys_core_pd(keys, X)
        else:
            values = keys_core_pd_v2(load_keys, X)

        values = values.groupby(as_index=False,
                                level='stock_code',
                                sort=False).apply(lambda x:x.pct_change(params))

        # TODO - 将inf设置为nan
        values[np.isinf(values['f1'])] = np.nan

        del load_keys

        return values['f1'].values

    return gp_cw_change_rate_qoq

# 3
@wrap_non_picklable_objects
def def_keyfunc_ts_cw_change_rate(keys, key_by_path=True):
    def gp_ts_cw_change_rate(X, d):
        if key_by_path:
            # load key
            load_keys = np.load(keys,
                                mmap_mode='r')
        else:
            load_keys = keys.copy()

        if len(load_keys) != len(X):
            return X

        # params = d_map(d_type_exam(d), base=base)
        params = d_map_limit(d_type_exam(d), base=2, lowlimit=cw_low_limit, uplimit=cw_up_limit)

        if not key_by_path:
            values = keys_core_pd(keys, X)
        else:
            values = keys_core_pd_v2(load_keys, X)

        values = values.groupby(as_index=False,
                                level='stock_code',
                                sort=False).apply(lambda x:x.pct_change(params))

        # TODO - 将inf设置为nan
        values[np.isinf(values['f1'])] = np.nan

        del load_keys

        return values['f1'].values

    return gp_ts_cw_change_rate

# 4
@wrap_non_picklable_objects
def def_keyfunc_cw_delta_qoq(keys, key_by_path=True):
    def gp_cw_delta_qoq(X):
        if key_by_path:
            # load key
            load_keys = np.load(keys,
                                mmap_mode='r')
        else:
            load_keys = keys.copy()

        if len(load_keys) != len(X):
            return X

        params = 1

        if not key_by_path:
            values = keys_core_pd(keys, X)
        else:
            values = keys_core_pd_v2(load_keys, X)

        values = values.groupby(as_index=False,
                                level='stock_code',
                                sort=False).apply(lambda x:x.diff(params))

        # TODO - 将inf设置为nan
        values[np.isinf(values['f1'])] = np.nan

        del load_keys

        return values['f1'].values

    return gp_cw_delta_qoq

# 5
@wrap_non_picklable_objects
def def_keyfunc_ts_cw_delta(keys, key_by_path=True):
    def gp_ts_cw_delta(X, d):
        if key_by_path:
            # load key
            load_keys = np.load(keys,
                                mmap_mode='r')
        else:
            load_keys = keys.copy()

        if len(load_keys) != len(X):
            return X

        # params = d_map(d_type_exam(d), base=1)
        params = d_map_limit(d_type_exam(d), base=2, lowlimit=cw_low_limit, uplimit=cw_up_limit)

        if not key_by_path:
            values = keys_core_pd(keys, X)
        else:
            values = keys_core_pd_v2(load_keys, X)

        values = values.groupby(as_index=False,
                                level='stock_code',
                                sort=False).apply(lambda x:x.diff(params))

        # TODO - 将inf设置为nan
        values[np.isinf(values['f1'])] = np.nan

        del load_keys

        return values['f1'].values

    return gp_ts_cw_delta

# 6
@wrap_non_picklable_objects
def def_keyfunc_ts_cw_bias(keys, key_by_path=True):
    def gp_ts_cw_bias(X, d):
        if key_by_path:
            # load key
            load_keys = np.load(keys,
                                mmap_mode='r')
        else:
            load_keys = keys.copy()

        if len(load_keys) != len(X):
            return X

        # params = d_map_limit(d_type_exam(d), base=4, limit=3)
        params = d_map_limit(d_type_exam(d), base=2, lowlimit=cw_low_limit, uplimit=cw_up_limit)

        if not key_by_path:
            values = keys_core_pd(keys, X)
        else:
            values = keys_core_pd_v2(load_keys, X)

        values = values.groupby(as_index=False,
                                level='stock_code',
                                sort=False).apply(lambda x:(x / x.rolling(window=params,
                                                                            min_periods=params).mean() - 1))

        # TODO - 将inf设置为nan
        values[np.isinf(values['f1'])] = np.nan

        del load_keys

        return values['f1'].values

    return gp_ts_cw_bias


'''linear regression at time series direction'''
# 7
@wrap_non_picklable_objects
def def_keyfunc_ts_cw_ls(keys, key_by_path=True):
    def gp_ts_cw_ls(X, d):
        if key_by_path:
            # load key
            load_keys = np.load(keys,
                                mmap_mode='r')
        else:
            load_keys = keys.copy()

        if len(load_keys) != len(X):
            return X

        params = d_map_limit(d_type_exam(d), base=2, lowlimit=cw_low_limit, uplimit=cw_up_limit)

        if not key_by_path:
            values = keys_core_pd(keys, X)
        else:
            values = keys_core_pd_v2(load_keys, X)

        values = values.astype(np.float64)

        def get_tech_analsis(df, params):
            df = df.copy()

            try:
                df['f1'] = talib.LINEARREG(df['f1'].values, params)
            except:
                df['f1'] = np.nan

            return df[['f1']]

        ls_values = values.groupby(as_index=False,
                                   level='stock_code',
                                   sort=False).apply(get_tech_analsis, params)

        # residual
        values = values - ls_values

        # TODO - 将inf设置为nan
        values[np.isinf(values['f1'])] = np.nan

        values = values.astype(np.float32)

        del load_keys

        return values['f1'].values

    return gp_ts_cw_ls


## linear regression slope
# 8
@wrap_non_picklable_objects
def def_keyfunc_ts_cw_ls_slope(keys, key_by_path=True):
    def gp_ts_cw_ls_slope(X, d):
        if key_by_path:
            # load key
            load_keys = np.load(keys,
                                mmap_mode='r')
        else:
            load_keys = keys.copy()

        if len(load_keys) != len(X):
            return X

        params = d_map_limit(d_type_exam(d), base=2, lowlimit=cw_low_limit, uplimit=cw_up_limit)

        if not key_by_path:
            values = keys_core_pd(keys, X)
        else:
            values = keys_core_pd_v2(load_keys, X)

        values = values.astype(np.float64)

        def get_tech_analsis(df, params):
            df = df.copy()

            try:
                df['f1'] = talib.LINEARREG_SLOPE(df['f1'].values, params)
                df['f1'] = talib.LINEARREG(df['f1'].values, params)

            except:
                df['f1'] = np.nan

            return df[['f1']]

        values = values.groupby(as_index=False,
                                level='stock_code',
                                sort=False).apply(get_tech_analsis, params)

        # TODO - 将inf设置为nan
        values[np.isinf(values['f1'])] = np.nan

        values = values.astype(np.float32)

        del load_keys

        return values['f1'].values

    return gp_ts_cw_ls_slope

## linear regression angle
# 9
@wrap_non_picklable_objects
def def_keyfunc_ts_cw_ls_angle(keys, key_by_path=True):
    def gp_ts_cw_ls_angle(X, d):
        if key_by_path:
            # load key
            load_keys = np.load(keys,
                                mmap_mode='r')
        else:
            load_keys = keys.copy()

        if len(load_keys) != len(X):
            return X

        params = d_map_limit(d_type_exam(d), base=2, lowlimit=cw_low_limit, uplimit=cw_up_limit)

        if not key_by_path:
            values = keys_core_pd(keys, X)
        else:
            values = keys_core_pd_v2(load_keys, X)

        values = values.astype(np.float64)

        def get_tech_analsis(df, params):
            df = df.copy()

            try:
                df['f1'] = talib.LINEARREG_ANGLE(df['f1'].values, params)
            except:
                df['f1'] = np.nan

            return df[['f1']]

        values = values.groupby(as_index=False,
                                level='stock_code',
                                sort=False).apply(get_tech_analsis, params)

        # TODO - 将inf设置为nan
        values[np.isinf(values['f1'])] = np.nan

        values = values.astype(np.float32)

        del load_keys

        return values['f1'].values

    return gp_ts_cw_ls_angle

## linear regression intercept
# 10
@wrap_non_picklable_objects
def def_keyfunc_ts_cw_ls_intercept(keys, key_by_path=True):
    def gp_ts_cw_ls_intercept(X, d):
        if key_by_path:
            # load key
            load_keys = np.load(keys,
                                mmap_mode='r')
        else:
            load_keys = keys.copy()

        if len(load_keys) != len(X):
            return X

        params = d_map_limit(d_type_exam(d), base=2, lowlimit=cw_low_limit, uplimit=cw_up_limit)

        if not key_by_path:
            values = keys_core_pd(keys, X)
        else:
            values = keys_core_pd_v2(load_keys, X)

        values = values.astype(np.float64)

        def get_tech_analsis(df, params):
            df = df.copy()

            try:
                df['f1'] = talib.LINEARREG_INTERCEPT(df['f1'].values, params)
            except:
                df['f1'] = np.nan

            return df[['f1']]

        values = values.groupby(as_index=False,
                                level='stock_code',
                                sort=False).apply(get_tech_analsis, params)

        # TODO - 将inf设置为nan
        values[np.isinf(values['f1'])] = np.nan

        values = values.astype(np.float32)

        del load_keys

        return values['f1'].values

    return gp_ts_cw_ls_intercept



